//
// Created by vivi on 15/09/2017.
//

#include <array>
#include "HairCore.h"

namespace XRwy {
	using namespace std;
	using namespace Eigen;

	HairSelleSimulator::HairSelleSimulator(const HairReader &reader, const HairSelleSimulatorConfiguration &conf):
			HairBaseSimulator(new HairSelleSimulatorConfiguration(conf)) {

		const auto & strands = reader.get_strands();
		const size_t strand_size = reader.get_strand_count();

		posv = vector<Vector3f>();
		velv = vector<Vector3f>();
		vector<bool> is_visiblev = vector<bool>();

		vector<size_t> total_particle_in_strand = vector<size_t>(strand_size, 0);

		for (size_t k = 0; k < strand_size; ++k) {

			const auto & strand = strands[k];

			for (size_t i = 0; i < strand.size(); ++i) {
				posv.push_back(strand[i]);
				velv.emplace_back(Vector3f::Zero());
				is_visiblev.push_back(true);
				++total_particle_in_strand[k];

				//we try to figure out whether we should insert an invisible particle between the local
				//particle i and i + 1

				//guide_line is used to ensure that two inserted particle should not in the same plane
				//to ensure this, when we inserted a virtual particle between particle i and particle i + 1,
				//we set the guide_line to new_inserted_particle - particle i. In that way, we only need to find
				//the normal vector for particle i, new_inserted_particle and particle i + 1, and insert the particle
				//in such direction
				Vector3f guide_line = Vector3f::Random();

				if (i != strand.size() - 1) {
					const auto & line = strand[i + 1] - strand[i];
					const float line_length = line.norm();

					//if the line is colinear with the before one or the after one, we should insert virtual particle
					if ((i >= 1 && HairCoreMath::is_colinear(strand[i] - strand[i - 1], line))
					    || (i + 2 < strand.size() && HairCoreMath::is_colinear(line, strand[i + 2] - strand[i + 1]))) {

						//we must ensure the guide_line cross with the line should not yield 0 vector
						while ((guide_line = guide_line.cross(line)).norm() < 1e-9)
							guide_line = Vector3f::Random();

						const Vector3f normal_line = guide_line.cross(line).normalized();
						const Vector3f inserted_particle = strand[i] + 0.5f * line + HairCoreMath::HALF_SQRT_3 * line_length * normal_line;

						posv.push_back(inserted_particle);
						velv.emplace_back(Vector3f::Zero());
						is_visiblev.push_back(false);
						++total_particle_in_strand[k];

						guide_line = normal_line; //use this to avoid most of random assignment for the guide line
					}
				}
			}
		}

		assert(posv.size() == velv.size());

		//now all particles have been pushed, not extra memory re-allocation will happen
		//we can now push particles to particle_view
		vector<ParticleView> views;
		size_t current_strand_to_push_idx = 0;

		for (size_t i = 0; i < posv.size(); ++i) {

			views.emplace_back(&posv[i], &velv[i], is_visiblev[i]);

			//update the push idx to next strand
			if (views.size() == total_particle_in_strand[current_strand_to_push_idx]) {
				push_strand(views);
				views.clear();
				++current_strand_to_push_idx;
			}
		}

		const size_t particle_size = particles.size();

		for (size_t i = 0; i < particle_size; ++i) {

			const size_t current_strand_id = particles[i].strand_id;

			//link stretch spring
			if (i + 1 < particle_size && particles[i + 1].strand_id == current_strand_id)
				springs.push_back(HairSelleSimulatorSpring{&this->particles[i], &this->particles[i + 1], conf.k, conf.damping});

			//link bending spring
			if (i + 2 < particle_size && particles[i + 2].strand_id == current_strand_id)
				springs.push_back(HairSelleSimulatorSpring{&this->particles[i], &this->particles[i + 2], conf.k, conf.damping});

			//link torison spring
			if (i + 3 < particle_size && particles[i + 3].strand_id == current_strand_id)
				springs.push_back(HairSelleSimulatorSpring{&this->particles[i], &this->particles[i + 3], conf.k, conf.damping});
		}
	}


	void HairSelleSimulator::strand_dynamic(const Affine3f & transformation, const float t) {

//#ifndef NDEBUG
//		cout << "-----------------------------------" << endl;
//		cout << "Before particle info:" << endl;
//		for (size_t i = 0; i < posv.size(); ++i)
//			cout << i << " pos: " << Vector3f_to_string(posv[i]) <<  "vel: " << Vector3f_to_string(velv[i]) << endl;
//		cout << "Before springs info:" << endl;
//		for (size_t i = 0; i < springs.size(); ++i) {
//			const auto & spring = springs[i];
//			cout << i << "\n\tparticle 1:(pos=" << Vector3f_to_string(spring.p1->pos()) << ", vel=" << Vector3f_to_string(spring.p1->vel()) << ")\n\tparticle 2:(pos="
//			     << Vector3f_to_string(spring.p2->pos()) << ", vel=" << Vector3f_to_string(spring.p2->vel()) << ")\n\tcurrent length:" << (spring.get_pos2() - spring.get_pos1()).norm()
//			     << "\n\torigin length:" << spring.l0 << "\n\timpluse:" << Vector3f_to_string(spring.impulse(t)) <<  "\n\tmatrix:" << '\n' << spring.matrix(t) << endl;
//		}
//#endif

		const auto & conf = get_conf();

		const size_t dimension_count = get_particle_count() * 3;

		//just by solving the equation: A * dv = b
		SparseMatrix<float> A(dimension_count, dimension_count);
		vector<Triplet<float>> A_triplets;
		ConjugateGradient<SparseMatrix<float>, Lower | Upper> cg_solver;

		VectorXf b = VectorXf::Zero(dimension_count);
		VectorXf dv;

		//every spring yields 36 values float
		//the damping and the identity matrix takes get_particle_count() * 3 values
		A_triplets.reserve(dimension_count + springs.size() * 36);
		for (size_t i = 0; i < get_particle_count() * 3; ++i)
			A_triplets.push_back(Triplet<float>(i, i, 1 + conf.air_damping * t));

		//apply the b for external acceleration impulse
		for (size_t i = 0, vi = 0 ; i < dimension_count; i += 3, vi++)
			b.segment<3>(i) = velv[vi] + t * conf.acc;

		for (const auto & spring : springs) {
			Matrix3f dm = spring.matrix(t);

			//apply the b for spring accleration impulse
			Vector3f impulse = spring.impulse(t);
			b.segment<3>(spring.p1->particle_id * 3) += impulse;
			b.segment<3>(spring.p2->particle_id * 3) -= impulse;

			//apply the A for matrix
			size_t ids[2] = {spring.p1->particle_id, spring.p2->particle_id};
			for (const size_t id1 : ids)
				for (const size_t id2 : ids)
					for (size_t i = 0; i < 3; ++i)
						for (size_t j = 0; j < 3; ++j)
							A_triplets.push_back(Triplet<float>(id1 * 3 + i, id2 * 3 + j, (id1 == id2 ? 1.0f : -1.0f) * dm(i, j)));
		}

		A.setFromTriplets(A_triplets.begin(), A_triplets.end());

		cg_solver.compute(A);
		dv = cg_solver.solve(b);

		//now add the new delta velocity to the previous one
		//and try changing the position as well

		for (size_t i = 0, vi = 0; vi < velv.size(); ++vi, i += 3) {
			if (particles[vi].local_id != 0) {
				velv[vi] = dv.segment<3>(i);
				posv[vi] += t * velv[vi];
			}
			else {
				Vector3f root_pos = transformation * particles[vi].rest_pos;
				velv[vi] = (root_pos - posv[vi]) / t;
				posv[vi] = root_pos;
			}
		}

#ifndef NDEBUG
		cout << "A:\n" << A << endl;
		cout << "dv:\n" << dv << endl;
		cout << "b:\n" << b << endl;
		cout << "After particle info: " << endl;
		for (size_t i = 0; i < posv.size(); ++i)
			cout << i << "pos: (" << posv[i](0) << ',' << posv[i](1) << ',' << posv[i](2) << ") vel: (" << velv[i](0) << ',' << velv[i](1) << ',' << velv[i](2) << ')' << endl;
		cout << "-----------------------------------" << endl;
#endif
	}
}